{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aberration Transfer Functions\n",
    "\n",
    "One may find reference on the internet to an \"Aberration Transfer Function\" introduced by Shannon used to model the MTF of an aberrated imaging system as:\n",
    "\n",
    "\\begin{align*}\n",
    "DTF(\\nu) &= \\frac{2}{\\pi}\\left[\\arccos{\\nu} - \\nu \\sqrt{1 - \\nu^2}\\right] \\\\\n",
    "ATF(\\nu) &= 1 - \\left(\\frac{W_\\text{rms}}{0.18}\\right)^2\\left(1 - 4(\\nu - 0.5)^2\\right) \\\\\n",
    "MTF(\\nu) &= DTF(\\nu) \\times ATF(\\nu)\n",
    "\\end{align*}\n",
    "\n",
    "where $DTF$ is the diffraction-limited MTF, $ATF$ is the \"Aberration Transfer Function,\" $MTF$ is the modulation transfer function, and $W_\\text{rms}$ is the RMS wavefront error.\n",
    "\n",
    "In this example, we will show that this treatment should not be used if accuracy is desired from a model.  The example should also highlight the terse nature of examples such as this when calculated using prysm.\n",
    "\n",
    "We begin by importing some classes and functions from the library, and defining the ATF function as Shannon describes it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from prysm.otf import diffraction_limited_mtf\n",
    "from prysm import FringeZernike, PSF, MTF\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "def shannon_atf(nu, Wrms):\n",
    "    return 1 - ((Wrms / 0.18) ** 2 * (1 - 4 * (nu - 0.5) ** 2 ))\n",
    "\n",
    "%matplotlib inline\n",
    "plt.style.use('bmh')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Wrms_vals = [0.025, 0.05, 0.075, 0.1, 0.125]\n",
    "nu = np.linspace(0, 1, 100)\n",
    "atf_curves = []\n",
    "for Wrms in Wrms_vals:\n",
    "    atf_curves.append(shannon_atf(nu=nu, Wrms=Wrms))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for (curve, label) in zip(atf_curves, Wrms_vals):\n",
    "    ax.plot(nu, curve, label=label)\n",
    "\n",
    "ax.legend(title=r'RMS WFE [$\\lambda$]')\n",
    "ax.set(xlim=(0,1), xlabel='Normalized Spatial Frequency [a.u.]',\n",
    "       ylim=(0,1), ylabel='ATF [Rel. 1.0]',\n",
    "       title=\"Shannon's ATF, various wavefront errors\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll pick a few different Zernike modes and show the numerically derived version, generated by calculating the MTF numerically and dividing by the diffraction limited MTF for a circular aperture, given above as $DTF$.  The accuracy of prysm's MTF calculations is sufficiently high that we can ignore that as a reason for the discrepancy.  [The accuracy of prysm's MTF calculations is that we can ignore that as a reason for any discrepancy.](./Diffraction%20Limited%20PSF%20and%20MTF.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only 20 lines of code, half of which is looping or plotting!\n",
    "def render_atf_curves_zernike_mode(mode_index, Wrms_vals):\n",
    "    kwarg = {}\n",
    "    \n",
    "    real_mtfs = []\n",
    "    for Wrms in Wrms_vals:\n",
    "        kwarg[f'Z{mode_index}'] = Wrms\n",
    "        pupil = FringeZernike(**kwarg, norm=True, z_unit='waves')  # waves is the default, not really needed\n",
    "        psf = PSF.from_pupil(pupil, efl=2)  # normalized frequency makes this choice arbitrary\n",
    "        mtf = MTF.from_psf(psf)\n",
    "        \n",
    "        u, mtf_ = mtf.slices().x\n",
    "        real_mtfs.append(mtf_)\n",
    "    \n",
    "    cutoff = 1 / (psf.wavelength * psf.fno) * 1e3 # 1e3 is cy/um => cy/mm\n",
    "    normalized_frequencies = u / cutoff\n",
    "    diffraction_limit = diffraction_limited_mtf(psf.fno, psf.wavelength, frequencies=u)\n",
    "    \n",
    "    # don't plot quite all of the curve, division by almost zero is a problem at the end\n",
    "    fig, ax = plt.subplots()\n",
    "    for (curve, label) in zip(real_mtfs, Wrms_vals):\n",
    "        atf = curve / diffraction_limit \n",
    "        ax.plot(normalized_frequencies[:-5], atf[:-5], label=label)\n",
    "    \n",
    "    ax.legend(title=f'RMS WFE [$\\lambda$]')\n",
    "    ax.set(xlim=(0,1), xlabel='Normalized Spatial Frequency',\n",
    "           ylim=(0,1), ylabel='ATF',\n",
    "           title=f'Numerically derived ATF for Fringe Zernike term Z{mode_index}')\n",
    "    \n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Z4 = defocus, the lowest-order Zernike error to affect imaging (and MTF)\n",
    "render_atf_curves_zernike_mode(4, Wrms_vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The curve looks broadly similar, but the belly reaches down quite a bit further.  What about higher order terms?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Z9 = \"zernike primary spherical\" -- low-order spherical aberration\n",
    "render_atf_curves_zernike_mode(9, Wrms_vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that for _low order_ spherical aberration, the curves look very different.  What if we had a higher order variant?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Z 25 = \"zernike tertiary spherical\" -- 8th order spherical aberration, in Hopkins' wave aberration expansion\n",
    "render_atf_curves_zernike_mode(25, Wrms_vals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even worse.  These are lots of squiggly lines, what if we directly compare a real ATF for a reasonable wavefront vs Shannon's ATF equation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# most of this code is just copy pasted from above\n",
    "pupil = FringeZernike(Z9=0.1, norm=True, z_unit='waves')  # waves is the default, not really needed\n",
    "psf = PSF.from_pupil(pupil, efl=2)  # normalized frequency makes this choice arbitrary\n",
    "mtf = MTF.from_psf(psf)\n",
    "\n",
    "u, mtf_ = mtf.slices().x\n",
    "\n",
    "diffraction_limit = diffraction_limited_mtf(psf.fno, psf.wavelength, frequencies=u)\n",
    "\n",
    "real_atf = mtf_ / diffraction_limit\n",
    "unormalized = u / (1 / (psf.wavelength * psf.fno) * 1e3)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(nu, shannon_atf(nu, 0.1), label=\"Shannon's eq.\")\n",
    "ax.plot(unormalized[:-5], real_atf[:-5], label='Numerical Solution')\n",
    "\n",
    "ax.legend(title='Method')\n",
    "ax.set(xlim=(0,1), xlabel='Normalized Spatial Frequency',\n",
    "       ylim=(0,1), ylabel='ATF',\n",
    "       title=r'Z9, RMS WFE = 0.1 $\\lambda$');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not a good match.  What if we look at the peak error in Shannon's equation as a function of Zernike index and RMS WFE corresponding to the Marechal critera?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def render_atf_peakerror_vs_zernike(max_zernike=36, rms_wfe= 1 / 14): # 1 / 14 is the Marechal criteria\n",
    "    # a lot of this code is similar to the earlier function\n",
    "    peak_errors = []\n",
    "    \n",
    "    # calculate one pilot case to get the metadata for the diffraction limited MTF.  This is a performance optimization\n",
    "    pupil = FringeZernike()\n",
    "    psf = PSF.from_pupil(pupil, efl=2)\n",
    "    mtf = MTF.from_psf(psf)\n",
    "    u, t = mtf.slices().x\n",
    "    \n",
    "    diffraction_limit = diffraction_limited_mtf(psf.fno, psf.wavelength, frequencies=u)\n",
    "    cutoff = 1 / (psf.wavelength * psf.fno) * 1e3\n",
    "    normalized_frequencies = u / cutoff\n",
    "    \n",
    "    shannon = shannon_atf(normalized_frequencies, rms_wfe)\n",
    "    \n",
    "    idxs = list(range(max_zernike+1))\n",
    "    for i in idxs:\n",
    "        kwarg = {}\n",
    "        kwarg[f'Z{i+1}'] = rms_wfe\n",
    "        pupil = FringeZernike(**kwarg, norm=True, z_unit='waves')  # waves is the default, not really needed\n",
    "        psf = PSF.from_pupil(pupil, efl=2)  # normalized frequency makes this choice arbitrary\n",
    "        mtf = MTF.from_psf(psf)\n",
    "        \n",
    "        cutoff = 1 / (psf.wavelength * psf.fno) * 1e3 # 1e3 is cy/um => cy/mm\n",
    "        u, mtf_ = mtf.slices().x\n",
    "        \n",
    "        atf = mtf_ / diffraction_limit\n",
    "        difference = abs(shannon[:-10] - atf[:-10])  # erode a little more of the end here for high order cases\n",
    "        peak_errors.append(difference.max())\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(idxs, peak_errors)\n",
    "    ax.set(xlim=(0,max_zernike), xlabel=\"Wavefront's Zernike index\",\n",
    "           ylim=(0,1), ylabel=\"Peak error of Shannon's ATF equation\",\n",
    "           title=f'RMS WFE = {rms_wfe:.3f}' + r'$\\lambda$')\n",
    "    \n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "render_atf_peakerror_vs_zernike(36, rms_wfe=1 / 14)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a wavefront at the Marechal criteria, the error can be as high as 0.7.  Since MTF must be within the range [0, 1], this means the error is at least 70%.  What if we had a tenth wave RMS?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "render_atf_peakerror_vs_zernike(36, rms_wfe=1 / 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the absolute error can be as high as 0.9, again at least 90% due to the normalization of MTF.\n",
    "\n",
    "Since these errors are so large, we can conclude that Shannon's ATF function should not be used if accuracy is desired."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}